{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Head Attention\n",
    "\n",
    "Multi-Head Attention is an important part of all Transformer-based models.\n",
    "This tutorial will show how to write it and how to then optimize it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-02 10:31:43.775 | DEBUG    | ttnn:<module>:83 - Initial ttnn.CONFIG:\n",
      "Config{cache_path=/home/ubuntu/.cache/ttnn,model_cache_path=/home/ubuntu/.cache/ttnn/models,tmp_dir=/tmp/ttnn,enable_model_cache=false,enable_fast_runtime_mode=true,throw_exception_on_fallback=false,enable_logging=false,enable_graph_report=false,enable_detailed_buffer_report=false,enable_detailed_tensor_report=false,enable_comparison_mode=false,comparison_mode_should_raise_exception=false,comparison_mode_pcc=0.9999,root_report_path=generated/ttnn/reports,report_name=std::nullopt,std::nullopt}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-07-02 10:31:43.820 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.820 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.833 | info     |          Device | Opening user mode device driver (tt_cluster.cpp:174)\n",
      "2025-07-02 10:31:43.833 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.833 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.838 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.838 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.843 | info     |   SiliconDriver | Harvesting mask for chip 0 is 0x40 (NOC0: 0x40, simulated harvesting mask: 0x0). (cluster.cpp:282)\n",
      "2025-07-02 10:31:43.874 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 10:31:43.881 | info     |   SiliconDriver | Opening local chip ids/pci ids: {0}/[0] and remote chip ids {} (cluster.cpp:147)\n",
      "2025-07-02 10:31:43.885 | info     |   SiliconDriver | Software version 6.0.0, Ethernet FW version 6.15.0 (Device 0) (cluster.cpp:1039)\n",
      "2025-07-02 10:31:43.900 | info     |           Metal | AI CLK for device 0 is:   1000 MHz (metal_context.cpp:122)\n",
      "2025-07-02 10:31:44.182 | info     |           Metal | Initializing device 0. Program cache is enabled (device.cpp:429)\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import torch\n",
    "import ttnn\n",
    "\n",
    "torch.manual_seed(0)\n",
    "\n",
    "device_id = 0\n",
    "dispatch_core_type = ttnn.device.DispatchCoreType.ETH\n",
    "if os.environ.get(\"ARCH_NAME\") and \"grayskull\" in os.environ.get(\"ARCH_NAME\"):\n",
    "    dispatch_core_type = ttnn.device.DispatchCoreType.WORKER\n",
    "device = ttnn.open_device(device_id=device_id, l1_small_size=8192, dispatch_core_config=ttnn.device.DispatchCoreConfig(dispatch_core_type))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enable program cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-07-02 10:31:44.774 | info     |           Metal | Enabling program cache on MeshDevice 1 (mesh_device.cpp:537)\n"
     ]
    }
   ],
   "source": [
    "device.enable_program_cache()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write Multi-Head Attention using ttnn\n",
    "\n",
    "Multi-head can be implemented in `torch` using just 6 operations:\n",
    "\n",
    "1. `torch.matmul`\n",
    "2. `torch.add`\n",
    "3. `torch.reshape`\n",
    "4. `torch.permute`\n",
    "5. `torch.mul`\n",
    "6. `torch.softmax`\n",
    "\n",
    "`ttnn` provides the exact same APIs to do that and therefore multi-head attention can be implemented in a very similar fashion. Except, when using `ttnn`, the user should be mindful of the tensor layout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    query_weight,\n",
    "    query_bias,\n",
    "    key_weight,\n",
    "    key_bias,\n",
    "    value_weight,\n",
    "    value_bias,\n",
    "    output_weight,\n",
    "    output_bias,\n",
    "    *,\n",
    "    num_heads,\n",
    "):\n",
    "    fallback_reshape = ttnn.get_fallback_function(ttnn.reshape) \n",
    "       \n",
    "    batch_size, sequence_size, hidden_size = hidden_states.shape\n",
    "    head_size = hidden_size // num_heads\n",
    "\n",
    "    query = hidden_states @ query_weight\n",
    "    query = query + query_bias\n",
    "    query = ttnn.to_layout(query, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "    query = fallback_reshape(query, (batch_size, sequence_size, num_heads, head_size))\n",
    "    query = ttnn.to_layout(query, layout=ttnn.TILE_LAYOUT)\n",
    "    query = ttnn.permute(query, (0, 2, 1, 3))\n",
    "\n",
    "    key = hidden_states @ key_weight\n",
    "    key = key + key_bias\n",
    "    key = ttnn.to_layout(key, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "    key = fallback_reshape(key, (batch_size, sequence_size, num_heads, head_size))\n",
    "    key = ttnn.to_layout(key, layout=ttnn.TILE_LAYOUT)\n",
    "    key = ttnn.permute(key, (0, 2, 3, 1))\n",
    "\n",
    "    value = hidden_states @ value_weight\n",
    "    value = value + value_bias\n",
    "    value = ttnn.to_layout(value, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "    value = fallback_reshape(value, (batch_size, sequence_size, num_heads, head_size))\n",
    "    value = ttnn.to_layout(value, layout=ttnn.TILE_LAYOUT)\n",
    "    value = ttnn.permute(value, (0, 2, 1, 3))\n",
    "\n",
    "    attention_scores = query @ key\n",
    "    attention_scores = attention_scores * (1 / (head_size**0.5))\n",
    "    attention_scores += attention_mask\n",
    "    attention_probs = ttnn.softmax(attention_scores, dim=-1)\n",
    "\n",
    "    context_layer = attention_probs @ value\n",
    "    context_layer = ttnn.permute(context_layer, (0, 2, 1, 3))\n",
    "    context_layer = ttnn.to_layout(context_layer, layout=ttnn.ROW_MAJOR_LAYOUT)\n",
    "    context_layer = fallback_reshape(context_layer, (batch_size, sequence_size, hidden_size))\n",
    "    context_layer = ttnn.to_layout(context_layer, layout=ttnn.TILE_LAYOUT)\n",
    "\n",
    "    self_output = context_layer @ output_weight\n",
    "    self_output = self_output + output_bias\n",
    "\n",
    "    return self_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the model is written, let's create input tensors to run it and test it"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "sequence_size = 384\n",
    "num_heads = 16\n",
    "head_size = 64\n",
    "hidden_size = num_heads * head_size"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize activations and weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_states = ttnn.rand((batch_size, sequence_size, hidden_size), dtype=ttnn.bfloat16, device=device)\n",
    "attention_mask = ttnn.rand((batch_size, 1, 1, sequence_size), dtype=ttnn.bfloat16, device=device)\n",
    "query_weight = ttnn.rand((hidden_size, hidden_size), dtype=ttnn.bfloat16, device=device)\n",
    "query_bias = ttnn.rand((hidden_size,), dtype=ttnn.bfloat16, device=device, memory_config=ttnn.L1_MEMORY_CONFIG)\n",
    "key_weight = ttnn.rand((hidden_size, hidden_size), dtype=ttnn.bfloat16, device=device)\n",
    "key_bias = ttnn.rand((hidden_size,), dtype=ttnn.bfloat16, device=device, memory_config=ttnn.L1_MEMORY_CONFIG)\n",
    "value_weight = ttnn.rand((hidden_size, hidden_size), dtype=ttnn.bfloat16, device=device)\n",
    "value_bias = ttnn.rand((hidden_size,), dtype=ttnn.bfloat16, device=device, memory_config=ttnn.L1_MEMORY_CONFIG)\n",
    "output_weight = ttnn.rand((hidden_size, hidden_size), dtype=ttnn.bfloat16, device=device)\n",
    "output_bias = ttnn.rand((hidden_size,), dtype=ttnn.bfloat16, device=device, memory_config=ttnn.L1_MEMORY_CONFIG)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the first iteration of Multi-Head Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    query_weight,\n",
    "    query_bias,\n",
    "    key_weight,\n",
    "    key_bias,\n",
    "    value_weight,\n",
    "    value_bias,\n",
    "    output_weight,\n",
    "    output_bias,\n",
    "    num_heads=num_heads,\n",
    ")\n",
    "end = time.time()\n",
    "duration = end - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Multi-head attention ran in 4.554651975631714 seconds for the first iteration\n"
     ]
    }
   ],
   "source": [
    "print(f\"Multi-head attention ran in {duration} seconds for the first iteration\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a subsequent iteration of Multi-Head Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "output = multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    query_weight,\n",
    "    query_bias,\n",
    "    key_weight,\n",
    "    key_bias,\n",
    "    value_weight,\n",
    "    value_bias,\n",
    "    output_weight,\n",
    "    output_bias,\n",
    "    num_heads=num_heads,\n",
    ")\n",
    "end = time.time()\n",
    "duration = end - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Multi-head attention ran in 0.05135488510131836 seconds for the subsequent iteration because of the program cache\n"
     ]
    }
   ],
   "source": [
    "print(f\"Multi-head attention ran in {duration} seconds for the subsequent iteration because of the program cache\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write optimized version of Multi-Head Attention\n",
    "\n",
    "Optimized version of the multi-head attention can be written by:\n",
    "\n",
    "- Tilizing all of the tensors ahead of time\n",
    "- Using more performant matmuls that fuse bias and specify the number of cores they execute on\n",
    "- Putting every tensor into L1\n",
    "- Using bfloat8_b data_type\n",
    "- Using custom `ttnn.transformer` operations instead of `ttnn.permute` and `ttnn.reshape`\n",
    "\n",
    "`ttnn.deallocate` calls are needed because otherwise, the cores on the device will run out of the L1 memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimized_multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    fused_qkv_weight,\n",
    "    fused_qkv_bias,\n",
    "    self_output_weight,\n",
    "    self_output_bias,\n",
    "    *,\n",
    "    num_heads,\n",
    "    num_cores_x=12,\n",
    "):\n",
    "    batch_size, _, hidden_size = hidden_states.shape\n",
    "    head_size = hidden_size // num_heads\n",
    "    \n",
    "    hidden_states = ttnn.to_layout(hidden_states, ttnn.TILE_LAYOUT)\n",
    "\n",
    "    fused_qkv_output = ttnn.linear(\n",
    "        hidden_states,\n",
    "        fused_qkv_weight,\n",
    "        bias=fused_qkv_bias,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "        dtype=ttnn.bfloat8_b,\n",
    "        core_grid=ttnn.CoreGrid(y=batch_size, x=num_cores_x),\n",
    "    )\n",
    "\n",
    "    (\n",
    "        query,\n",
    "        key,\n",
    "        value,\n",
    "    ) = ttnn.transformer.split_query_key_value_and_split_heads(\n",
    "        fused_qkv_output,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "        num_heads=num_heads,\n",
    "    )\n",
    "    ttnn.deallocate(fused_qkv_output)\n",
    "\n",
    "    attention_scores = ttnn.matmul(\n",
    "        query,\n",
    "        key,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "        dtype=ttnn.bfloat16,\n",
    "        core_grid=ttnn.CoreGrid(y=batch_size, x=num_cores_x),\n",
    "    )\n",
    "    ttnn.deallocate(query)\n",
    "    ttnn.deallocate(key)\n",
    "\n",
    "    attention_probs = ttnn.transformer.attention_softmax_(attention_scores, attention_mask=attention_mask, head_size=head_size)\n",
    "\n",
    "    context_layer = ttnn.matmul(\n",
    "        attention_probs,\n",
    "        value,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "        dtype=ttnn.bfloat8_b,\n",
    "        core_grid=ttnn.CoreGrid(y=batch_size, x=num_cores_x),\n",
    "    )\n",
    "    ttnn.deallocate(attention_probs)\n",
    "\n",
    "    context_layer_after_concatenate_heads = ttnn.transformer.concatenate_heads(\n",
    "        context_layer,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "    )\n",
    "    ttnn.deallocate(context_layer)\n",
    "\n",
    "    self_output = ttnn.linear(\n",
    "        context_layer_after_concatenate_heads,\n",
    "        self_output_weight,\n",
    "        bias=self_output_bias,\n",
    "        memory_config=ttnn.L1_MEMORY_CONFIG,\n",
    "        dtype=ttnn.bfloat16,\n",
    "        core_grid=ttnn.CoreGrid(y=batch_size, x=num_cores_x),\n",
    "    )\n",
    "    ttnn.deallocate(context_layer_after_concatenate_heads)\n",
    "\n",
    "    return self_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-process the parameters of the optimized model\n",
    "\n",
    "1. Fuse QKV weights and biases\n",
    "2. Reshape and tilize for the optimized operations using preprocess_linear_weight and preprocess_linear_bias\n",
    "3. Move to device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "qkv_weight = ttnn.concat([query_weight, key_weight, value_weight], dim=-1)\n",
    "qkv_weight = ttnn.to_device(ttnn.to_layout(qkv_weight, layout=ttnn.TILE_LAYOUT), device=device)\n",
    "\n",
    "qkv_bias = ttnn.from_torch(torch.cat([ttnn.to_torch(query_bias), ttnn.to_torch(key_bias), ttnn.to_torch(value_bias)], dim=-1), device=device)\n",
    "qkv_bias = ttnn.reshape(qkv_bias, (1, -1))\n",
    "qkv_bias = ttnn.to_layout(qkv_bias, layout=ttnn.TILE_LAYOUT)\n",
    "\n",
    "output_weight = ttnn.to_layout(output_weight, layout=ttnn.TILE_LAYOUT)\n",
    "\n",
    "output_bias = ttnn.reshape(output_bias, (1, -1))\n",
    "output_bias = ttnn.to_layout(output_bias, layout=ttnn.TILE_LAYOUT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the first iteration of the optimized Multi-Head Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "hidden_states = ttnn.to_layout(hidden_states, ttnn.TILE_LAYOUT)\n",
    "optimized_output = optimized_multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    qkv_weight,\n",
    "    qkv_bias,\n",
    "    output_weight,\n",
    "    output_bias,\n",
    "    num_heads=num_heads,\n",
    ")\n",
    "end = time.time()\n",
    "duration = end - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimized multi-head attention ran in 2.4207022190093994 seconds for the first iteration\n"
     ]
    }
   ],
   "source": [
    "print(f\"Optimized multi-head attention ran in {duration} seconds for the first iteration\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a subsequent iteration of the optimized Multi-Head Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "optimized_output = optimized_multi_head_attention(\n",
    "    hidden_states,\n",
    "    attention_mask,\n",
    "    qkv_weight,\n",
    "    qkv_bias,\n",
    "    output_weight,\n",
    "    output_bias,\n",
    "    num_heads=num_heads,\n",
    ")\n",
    "end = time.time()\n",
    "duration = end - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimized multi-head attention ran in 0.0037069320678710938 seconds for the subsequent iteration because of the program cache\n"
     ]
    }
   ],
   "source": [
    "print(f\"Optimized multi-head attention ran in {duration} seconds for the subsequent iteration because of the program cache\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the optimized multi-head attention is 2 orders of magnitude faster than the initial version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check that the output of the optimized version matches the output of the original implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch_output = ttnn.to_torch(output)\n",
    "torch_optimized_output = ttnn.to_torch(optimized_output)\n",
    "\n",
    "assert torch.allclose(torch_output, torch_optimized_output)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Close the device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-07-02 10:31:53.340 | info     |           Metal | Closing mesh device 1 (mesh_device.cpp:488)\n",
      "2025-07-02 10:31:53.341 | info     |           Metal | Closing mesh device 0 (mesh_device.cpp:488)\n",
      "2025-07-02 10:31:53.342 | info     |           Metal | Closing device 0 (device.cpp:469)\n",
      "2025-07-02 10:31:53.342 | info     |           Metal | Disabling and clearing program cache on device 0 (device.cpp:781)\n"
     ]
    }
   ],
   "source": [
    "ttnn.close_device(device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
